<style>
  body {
    margin: 0;
    font-family: Arial, sans-serif;
    overflow: hidden;
    background: #222;
  }

  svg {
    width: 100vw;
    height: 100vh;
    display: block;
  }

  .node {
    cursor: pointer;
  }

  .link {
    stroke-opacity: 0.7;
  }

  .node text {
    pointer-events: none;
    font-size: 10px;
    fill: #fff;
    font-weight: bold;
    text-shadow: 0 0 3px #000;
    opacity: 1;
  }
</style>

<script type="module">
  import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7/+esm";

  // Initialize the visualization
  let svg, simulation, link, node, width, height;
  let container; // Make container global
  let linksGroup, nodesGroup; // Add group references
  const nodes = [];
  const links = [];

  // Initial dataset
  const initialConcepts = {
    // "relevant_concepts": ["AI", "machine learning", "data science", "neural networks", "deep learning", "natural language processing"],
    // "orthogonal_concepts": ["academic research", "philosophy of mind", "cognitive science"],
    // "user_needs": ["understand AI", "learn machine learning basics"],
    // "assistant_needs": ["provide clear explanations", "avoid jargon"],
    // "strategies": ["use analogies", "break down complex concepts"]
  };

  // Color scheme for different types of nodes (Open WebUI inspired)
  const colorMap = {
    "relevant_concepts": "#dedede",   // Purple (primary brand color)
    "orthogonal_concepts": "#6366f1", // Indigo
    "user_needs": "#10b981",          // Emerald
    "assistant_needs": "#f59e0b",     // Amber
    "strategies": "#f97316"           // Orange
  };

  // Initialize the visualization
  function initializeVisualization() {
    // Set up SVG container
    svg = d3.select("body").append("svg")
      .attr("width", "100%")
      .attr("height", "100%");

    // Add zoom and pan capabilities
    const zoom = d3.zoom()
      .scaleExtent([0.1, 4])
      .on("zoom", (event) => {
        container.attr("transform", event.transform);
      });

    svg.call(zoom);

    // Add a container for all visualization elements
    container = svg.append("g")
      .attr("class", "container");

    // Create group for links and nodes and save references
    linksGroup = container.append("g").attr("class", "links");
    nodesGroup = container.append("g").attr("class", "nodes");

    // Create a gradient definition for links
    const defs = svg.append("defs");

    // Create a glow filter definition (moved to initialization to fix clipping)
    const glow = defs.append("filter")
      .attr("id", "glow")
      .attr("x", "-50%")
      .attr("y", "-50%")
      .attr("width", "200%")
      .attr("height", "200%");

    const glowStronger = defs.append("filter")
      .attr("id", "glowStronger")
      .attr("x", "-50%")
      .attr("y", "-50%")
      .attr("width", "200%")
      .attr("height", "200%");

    glow.append("feGaussianBlur")
      .attr("stdDeviation", "2")
      .attr("result", "coloredBlur");

    const feMerge = glow.append("feMerge");
    feMerge.append("feMergeNode").attr("in", "coloredBlur");
    feMerge.append("feMergeNode").attr("in", "SourceGraphic");


    glowStronger.append("feGaussianBlur")
      .attr("stdDeviation", "6")
      .attr("result", "coloredBlur");

    const feMergeStronger = glowStronger.append("feMerge");
    feMergeStronger.append("feMergeNode").attr("in", "coloredBlur");
    feMergeStronger.append("feMergeNode").attr("in", "SourceGraphic");

    // Get dimensions
    updateDimensions();

    // Create the simulation
    simulation = d3.forceSimulation()
      .force("link", d3.forceLink().id(d => d.id).distance(120)) // was 10
      .force("charge", d3.forceManyBody().strength(-250))
      .force("x", d3.forceX(width / 2).strength(0.03))
      .force("y", d3.forceY(height / 2).strength(0.03))
      .force("collide", d3.forceCollide().radius(32)); // was 50

    // Populate with initial data
    populateInitialData();

    // Update on window resize
    window.addEventListener("resize", () => {
      updateDimensions();
      simulation.force("center", d3.forceCenter(width / 2, height / 2));
      simulation.alpha(0.5).restart();
    });
  }

  // Update width and height
  function updateDimensions() {
    width = window.innerWidth;
    height = window.innerHeight;
    svg.attr("width", width).attr("height", height);
  }

  // Populate with initial data
  function populateInitialData() {
    // Add nodes for each category
    Object.entries(initialConcepts).forEach(([category, concepts]) => {
      concepts.forEach(concept => {
        addNode(concept, category);
      });
    });
  }

  // Add a new node to the graph
  function addNode(label, category = "relevant_concepts") {
    // Check if node already exists
    if (nodes.find(n => n.id === label)) {
      console.log(`Node ${label} already exists`);
      return;
    }

    // Add to nodes array
    const newNode = {
      id: label,
      label: label,
      category: category,
      x: (width / 2),
      y: (height / 2),
    };

    nodes.push(newNode);

    // Update the simulation
    updateVisualization();

    // Return the node for chaining
    return newNode;
  }

  // Track highlighted (linked) nodes
  const highlightedNodes = new Set();

  // Link two nodes by their labels
  function linkNodes(fromLabel, toLabel) {
    // Check if both nodes exist
    const fromNode = nodes.find(n => n.id === fromLabel);
    const toNode = nodes.find(n => n.id === toLabel);

    if (!fromNode || !toNode) {
      throw new Error(`Cannot link: one or both nodes don't exist (${fromLabel}, ${toLabel})`);
    }

    // Check if link already exists
    if (links.some(l => (l.source === fromLabel && l.target === toLabel) ||
      (l.source === toLabel && l.target === fromLabel))) {
      console.log(`Link between ${fromLabel} and ${toLabel} already exists`);
      return;
    }

    // Add to links array
    const newLink = {
      source: fromLabel,
      target: toLabel
    };

    links.push(newLink);

    // Mark both nodes as highlighted
    highlightedNodes.add(fromLabel);
    highlightedNodes.add(toLabel);

    // Update the visualization (will apply highlight)
    updateVisualization();

    // No need for highlightLinkedNodes anymore
    // Return the link for chaining
    return newLink;
  }

  // Update the visualization with current nodes and links
  function updateVisualization() {
    // Update links in linksGroup
    link = linksGroup.selectAll("line").data(links, d => `${d.source.id || d.source}-${d.target.id || d.target}`);
    link.exit().remove();
    const linkEnter = link
      .enter().append("line")
      .attr("class", "link")
      .attr("stroke", "#999")
      .attr("stroke-width", 2);
    link = linkEnter.merge(link);

    // Update nodes in nodesGroup
    node = nodesGroup.selectAll(".node").data(nodes, d => d.id);
    node.exit().remove();
    const nodeEnter = node
      .enter().append("g")
      .attr("class", "node")
      .call(drag(simulation));

    // Add rounded rectangle (pill) to node
    nodeEnter.append("rect")
      .attr("x", d => -d.label.length * 4 - 10)
      .attr("y", -12)
      .attr("width", d => d.label.length * 8 + 20)
      .attr("height", 24)
      .attr("rx", 12)
      .attr("ry", 12)
      .style("fill", d => colorMap[d.category] || "#6b7280")
      .style("stroke", "#e5e7eb")
      .style("stroke-width", 2)
      .style("filter", "url(#glow)")
      .attr("transform", "scale(0)") // Start with scale 0 for springy motion
      .transition() // Add springy scale-in effect
      .duration(600)
      .ease(d3.easeElastic.period(0.75))
      .attr("transform", "scale(1)");

    // Add text to node (centered in the pill)
    nodeEnter.append("text")
      .attr("dy", 4)
      .attr("text-anchor", "middle")
      .text(d => d.label)
      .style("fill", "#2d2d2d")
      .style("font-weight", "600")
      .style("font-size", "14px")
      .style("text-shadow", "0 1px 2px rgba(0,0,0,0.3)")
      .attr("transform", "scale(0)")
      .transition()
      .duration(600)
      .ease(d3.easeElastic.period(0.75))
      .attr("transform", "scale(1)");

    node = nodeEnter.merge(node);

    // Apply highlight to all highlighted nodes
    node.each(function(d) {
      const isHighlighted = highlightedNodes.has(d.id);
      d3.select(this).select('rect')
        .transition()
        .duration(300)
        .attr('transform', isHighlighted ? 'scale(1.25)' : 'scale(1)')
        .style('filter', isHighlighted ? 'url(#glowStronger)' : 'url(#glow)');
      d3.select(this).select('text')
        .transition()
        .duration(300)
        .attr('transform', isHighlighted ? 'scale(1.15)' : 'scale(1)')
    });

    // No need to re-select globally, keep selections scoped

    // Update simulation
    simulation.nodes(nodes);
    simulation.force("link").links(links);
    simulation.alpha(0.2).restart();
    // Always raise nodes group to top so nodes render above links
    nodesGroup.raise();
  }

  // Drag functions for nodes
  function drag(simulation) {
    function dragstarted(event) {
      if (!event.active) simulation.alphaTarget(0.3).restart();
      event.subject.fx = event.subject.x;
      event.subject.fy = event.subject.y;
    }

    function dragged(event) {
      event.subject.fx = event.x;
      event.subject.fy = event.y;
    }

    function dragended(event) {
      if (!event.active) simulation.alphaTarget(0);
      event.subject.fx = null;
      event.subject.fy = null;
    }

    return d3.drag()
      .on("start", dragstarted)
      .on("drag", dragged)
      .on("end", dragended);
  }

  // Update positions in the animation loop
  function ticked() {
    // Always select current DOM elements for links and nodes
    container.selectAll('.link')
      .attr('x1', d => d.source.x)
      .attr('y1', d => d.source.y)
      .attr('x2', d => d.target.x)
      .attr('y2', d => d.target.y);

    container.selectAll('.node')
      .attr('transform', d => `translate(${d.x}, ${d.y})`);
  }

  // Event handling setup
  const handlers = {
    "boost.listener.event": handleBoostEvent,
    "chat.completion.chunk": handleCompletionChunk,
    'boost.concept': function ({ data }) {
      addNode(data.concept, "relevant_concepts");
    },
    'boost.linked_concepts': function ({ data }) {
      const { concepts } = data;
      concepts.reduce((prev, curr) => {
        if (prev) {
          linkNodes(prev, curr);
        }
        return curr;
      }, null);
    },
  };

  function processChunk(chunk) {
    try {
      const data = JSON.parse(chunk.replace(/data: /, ""));
      const text = data.object;
      const handler = handlers[text];

      if (handler) {
        console.log("Processing chunk:", data);
        handler(data);
      }
    } catch (e) {
      console.error("Error processing chunk:", e);
    }
  }

  function handleCompletionChunk(chunk) {
    // addNode(getChunkContent(chunk), "relevant_concepts");
  }

  function handleBoostEvent(chunk) {
    const { event } = chunk;
    const handler = handlers[event];

    if (handler) {
      console.log("Handling boost event:", event);
      handler(chunk);
    } else {
      console.warn(`No handler for event: ${event}`);
    }
  }

  function getChunkContent(chunk) {
    return chunk.choices.map((choice) => choice.delta.content).join("\n");
  }

  // Initialize when document is loaded
  document.addEventListener("DOMContentLoaded", () => {
    initializeVisualization();
    simulation.on("tick", ticked);
    startListening();

    // Expose functions to window for debugging/demonstration
    window.addNode = addNode;
    window.linkNodes = linkNodes;
  });

  // Start listening for events
  async function startListening() {
    try {
      const listenerId = "<<listener_id>>";
      const boostUrl = '<<boost_public_url>>';

      const response = await fetch(
        `${boostUrl}/events/${listenerId}`,
        {
          headers: {
            Authorization: "Bearer sk-boost",
          },
        }
      );

      if (!response.ok) {
        throw new Error(`HTTP error! status: ${response.status}`);
      }

      const reader = response.body.getReader();

      while (true) {
        const { done, value } = await reader.read();
        if (done) {
          console.log("Stream complete");
          break;
        }

        try {
          const blob = new TextDecoder().decode(value);
          const chunks = blob.split("\n\n");

          for (const chunk of chunks) {
            if (chunk.trim()) {
              processChunk(chunk);
            }
          }
        } catch (e) {
          console.error("Error processing data:", e);
        }
      }
    } catch (error) {
      console.error("Error connecting to event stream:", error);
    }
  }
</script>